from vnpy_ctastrategy import (
    CtaTemplate,
    StopOrder,
    TickData,
    BarData,
    TradeData,
    OrderData,
    BarGenerator,
    ArrayManager,
)

from vnpy.trader.constant import Interval, Direction
from datetime import time, timedelta, datetime
import pandas as pd
import numpy as np
import tushare as ts


class FundamentalValueStrategy(CtaTemplate):
    """
    增强型基本面估值策略
    
    基于PE/PB相对行业均值的偏离度进行交易决策，并结合技术指标增加交易频率:
    1. 买入信号：PE或PB低于行业均值达到设定阈值，并且技术指标支持买入
    2. 卖出信号：PE或PB高于行业均值达到设定阈值，或者技术指标出现卖出信号
    3. 采用部分仓位管理，根据信号强度调整仓位大小
    """
    author = "VeighNa量化"

    # Tushare API令牌
    tushare_token = "2b98c7cf1077a21e7e16811e4486ae789c7ff675f91926159b876327"
    
    # 基础交易数量
    fixed_size = 100

    # 估值偏离阈值参数（放宽阈值，提高信号频率）
    pe_low_threshold = -0.05   # PE低于行业均值5%
    pe_high_threshold = 0.005  # PE高于行业均值0.5%
    pb_low_threshold = -0.05   # PB低于行业均值5%
    pb_high_threshold = 0.005  # PB高于行业均值0.5%
    
    # 仓位和更新周期参数
    max_pos = 1000              # 最大持仓量
    rebalance_days = 5         # 每5个交易日更新一次估值数据
    
    # 技术指标参数
    fast_ma_period = 5          # 快速均线周期
    slow_ma_period = 10         # 慢速均线周期
    rsi_period = 14             # RSI指标周期
    rsi_buy_threshold = 40      # RSI买入阈值
    rsi_sell_threshold = 60     # RSI卖出阈值
    
    # 止损参数
    stop_loss_pct = 0.05        # 止损比例
    
    # 策略参数
    parameters = [
        "tushare_token",
        "fixed_size",
        "pe_low_threshold", "pe_high_threshold",
        "pb_low_threshold", "pb_high_threshold",
        "max_pos", "rebalance_days",
        "fast_ma_period", "slow_ma_period",
        "rsi_period", "rsi_buy_threshold", "rsi_sell_threshold",
        "stop_loss_pct"
    ]

    # 策略变量
    variables = [
        "current_pe", "current_pb", 
        "industry_pe", "industry_pb",
        "pe_deviation", "pb_deviation",
        "days_since_update", "industry_code",
        "fast_ma", "slow_ma", "rsi_value"
    ]

    def __init__(self, cta_engine, strategy_name, vt_symbol, setting):
        """
        策略初始化
        """
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)
        
        # 数据管理
        self.bg = BarGenerator(self.on_bar)
        self.am = ArrayManager(100)
        
        # 估值指标初始化
        self.current_pe = 0.0
        self.current_pb = 0.0
        self.industry_pe = 0.0
        self.industry_pb = 0.0
        self.pe_deviation = 0.0
        self.pb_deviation = 0.0
        
        # 技术指标初始化
        self.fast_ma = 0.0
        self.slow_ma = 0.0
        self.rsi_value = 0.0
        
        # 行业信息与状态跟踪
        self.industry_code = ""
        self.days_since_update = 0
        
        # Tushare相关
        self.ts_api = None
        # 转换股票代码格式
        original_code = vt_symbol.split(".")[0]
        exchange = vt_symbol.split(".")[1]
        if exchange == "SSE":
            self.stock_code = f"{original_code}.SH"
        elif exchange == "SZSE":
            self.stock_code = f"{original_code}.SZ"
        else:
            self.stock_code = original_code
        
        # 交易管理
        self.entry_price = 0.0
        self.high_since_entry = 0.0  # 入场后的最高价
        self.position_stage = 0      # 0: 无仓位, 1: 部分建仓, 2: 满仓
        
    def _get_stock_industry(self):
        """
        获取股票所属行业
        """
        if not self.ts_api:
            return ""
        
        try:
            # 获取股票基本信息
            df = self.ts_api.stock_basic(
                ts_code=self.stock_code, 
                fields='ts_code,industry'
            )
            
            if not df.empty:
                # 获取行业代码
                self.industry_code = df['industry'].iloc[0]
                print(f"股票 {self.stock_code} 所属行业: {self.industry_code}")
                return self.industry_code
                
        except Exception as e:
            self.write_log(f"获取行业信息异常: {str(e)}")
        
        return ""
    
    def _update_valuation_data(self):
        """
        更新估值数据
        """
        if not self.ts_api:
            return
            
        try:
            print("\n更新估值数据:")
            print(f"股票代码: {self.stock_code}")
            
            # 获取最近的交易日期
            today = datetime.now()
            start_date = (today - timedelta(days=7)).strftime('%Y%m%d')
            end_date = today.strftime('%Y%m%d')
            
            try:
                # 获取交易日历
                df_trade_cal = self.ts_api.trade_cal(
                    exchange='SSE',
                    start_date=start_date,
                    end_date=end_date,
                    is_open='1'  # 1表示交易日
                )
                
                if df_trade_cal.empty:
                    print("错误: 无法获取交易日历")
                    return
                    
                # 获取最新的交易日期
                latest_trade_date = df_trade_cal['cal_date'].iloc[-1]
                print(f"最新交易日期: {latest_trade_date}")
                
                # 1. 获取股票当前估值
                print(f"\n1. 获取股票当前估值...")
                df_stock = self.ts_api.daily_basic(
                    ts_code=self.stock_code,
                    trade_date=latest_trade_date,
                    fields='ts_code,trade_date,pe,pb'
                )
                
                if df_stock.empty:
                    print(f"警告: 未获取到 {self.stock_code} 的估值数据")
                    return
                
                self.current_pe = float(df_stock['pe'].iloc[0])
                self.current_pb = float(df_stock['pb'].iloc[0])
                print(f"当前估值: PE={self.current_pe:.2f}, PB={self.current_pb:.2f}")
                
                # 获取行业代码
                if not self.industry_code:
                    self._get_stock_industry()
                
                if not self.industry_code:
                    print("警告: 无法获取行业信息，无法计算行业估值")
                    return
                
                # 2. 获取行业估值数据 - 优化版本
                print(f"\n2. 获取行业估值数据...")
                
                # 先获取同行业股票列表，最多取前50只（减轻API负担）
                df_industry_stocks = self.ts_api.stock_basic(
                    exchange='',
                    list_status='L',
                    fields='ts_code,industry,market',
                    industry=self.industry_code
                )
                
                if df_industry_stocks.empty:
                    print(f"警告: 未找到行业 {self.industry_code} 的股票")
                    return
                
                # 按照市值排序，取前50只大市值股票作为样本（更具代表性）
                df_industry_stocks = df_industry_stocks.sort_values('market', ascending=False)
                industry_samples = df_industry_stocks['ts_code'].head(30).tolist()
                
                # 分批获取这些股票的估值数据，每次最多10只
                batch_size = 10
                all_industry_data = []
                
                for i in range(0, len(industry_samples), batch_size):
                    batch = industry_samples[i:i+batch_size]
                    batch_str = ','.join(batch)
                    
                    df_batch = self.ts_api.daily_basic(
                        ts_code=batch_str,
                        trade_date=latest_trade_date,
                        fields='ts_code,trade_date,pe,pb'
                    )
                    
                    if not df_batch.empty:
                        all_industry_data.append(df_batch)
                
                if not all_industry_data:
                    print(f"警告: 未获取到行业估值样本数据")
                    return
                
                # 合并所有批次数据
                df_industry = pd.concat(all_industry_data)
                
                # 计算行业平均PE和PB (剔除异常值)
                pe_valid = df_industry[(df_industry['pe'] > 0) & (df_industry['pe'] < 300)]
                pb_valid = df_industry[(df_industry['pb'] > 0) & (df_industry['pb'] < 50)]
                
                if not pe_valid.empty:
                    self.industry_pe = pe_valid['pe'].mean()
                if not pb_valid.empty:
                    self.industry_pb = pb_valid['pb'].mean()
                
                print(f"行业平均估值: PE={self.industry_pe:.2f}, PB={self.industry_pb:.2f}")
                print(f"行业样本数量: {len(df_industry)}只股票")
                
                # 计算偏离度
                if self.industry_pe > 0 and self.current_pe > 0:
                    self.pe_deviation = (self.current_pe - self.industry_pe) / self.industry_pe
                
                if self.industry_pb > 0 and self.current_pb > 0:
                    self.pb_deviation = (self.current_pb - self.industry_pb) / self.industry_pb
                
                print(f"估值偏离度: PE偏离={self.pe_deviation:.2%}, PB偏离={self.pb_deviation:.2%}")
                
                # 重置更新天数计数
                self.days_since_update = 0
                
            except Exception as e:
                print(f"更新估值数据时出错: {str(e)}")
                print("错误详情:", e.__class__.__name__)
                
        except Exception as e:
            print(f"\n更新估值数据失败: {str(e)}")
    
    def on_init(self):
        """
        策略初始化，加载历史数据
        """
        self.write_log("策略初始化")
        
        # 初始化Tushare API
        if self.tushare_token:
            ts.set_token(self.tushare_token)
            self.ts_api = ts.pro_api()
            print("Tushare API初始化成功")
            
            # 获取估值数据
            self._update_valuation_data()
        
        self.load_bar(20)  # 加载更多历史数据以计算技术指标
        print("策略初始化完成")

    def on_start(self):
        """策略启动"""
        self.write_log("策略启动")
        
    def on_stop(self):
        """策略停止"""
        self.write_log("策略停止")

    def on_tick(self, tick: TickData):
        """Tick数据更新"""
        self.bg.update_tick(tick)

    def on_bar(self, bar: BarData):
        """K线数据更新"""
        print(f"\n收到K线数据: {bar.datetime}")
        print(f"价格: 开={bar.open_price:.2f}, 高={bar.high_price:.2f}, 低={bar.low_price:.2f}, 收={bar.close_price:.2f}")
        print(f"当前持仓: {self.pos}股")
        
        # 更新技术指标
        am = self.am
        am.update_bar(bar)
        
        if am.inited:
            # 计算技术指标
            self.fast_ma = am.sma(self.fast_ma_period)
            self.slow_ma = am.sma(self.slow_ma_period)
            self.rsi_value = am.rsi(self.rsi_period)
            
            print(f"技术指标: 快线MA{self.fast_ma_period}={self.fast_ma:.2f}, 慢线MA{self.slow_ma_period}={self.slow_ma:.2f}, RSI{self.rsi_period}={self.rsi_value:.2f}")
        
        # 更新最高价记录
        if self.pos > 0:
            self.high_since_entry = max(self.high_since_entry, bar.high_price)
        
        # 更新估值数据
        self.days_since_update += 1
        if self.days_since_update >= self.rebalance_days:
            print(f"\n达到再平衡周期({self.rebalance_days}天)，更新估值数据")
            self._update_valuation_data()
        
        # 交易逻辑
        if am.inited:
            # 检查止损条件
            if self.pos > 0 and self.entry_price > 0:
                # 固定止损
                if bar.close_price <= self.entry_price * (1 - self.stop_loss_pct):
                    print(f"\n触发止损: 当前价格({bar.close_price:.2f}) <= 止损价格({self.entry_price * (1 - self.stop_loss_pct):.2f})")
                    self.sell(bar.close_price, self.pos)
                    print(f"发出止损卖出委托: 价格={bar.close_price:.2f}, 数量={self.pos}")
                # 移动止损（从最高点回撤超过3%）
                elif bar.close_price <= self.high_since_entry * 0.97 and self.high_since_entry > self.entry_price * 1.05:
                    print(f"\n触发移动止损: 当前价格({bar.close_price:.2f}) <= 移动止损价格({self.high_since_entry * 0.97:.2f})")
                    print(f"最高价: {self.high_since_entry:.2f}, 入场价: {self.entry_price:.2f}")
                    self.sell(bar.close_price, self.pos)
                    print(f"发出移动止损卖出委托: 价格={bar.close_price:.2f}, 数量={self.pos}")
                else:
                    self._generate_trade_signal(bar)
            else:
                self._generate_trade_signal(bar)
        
        self.put_event()
    
    def _generate_trade_signal(self, bar: BarData):
        """
        生成交易信号 - 增强版本
        """
        # 确保有基本的估值数据和技术指标
        if not self.am.inited:
            return
            
        # 检查买入条件：基本面条件 OR 技术条件
        if self.pos == 0:
            # 基本面买入条件：PE或PB明显低于行业均值
            value_buy_signal = False
            if self.current_pe > 0 and self.industry_pe > 0 and self.current_pb > 0 and self.industry_pb > 0:
                value_buy_signal = (self.pe_deviation <= self.pe_low_threshold or 
                                   self.pb_deviation <= self.pb_low_threshold)
                                   
            # 技术面买入条件：金叉或RSI超卖
            tech_buy_signal = (self.fast_ma > self.slow_ma and 
                              self.rsi_value > 30 and 
                              self.rsi_value < 70)
            
            # 综合买入信号
            if tech_buy_signal and (value_buy_signal or self.rsi_value < self.rsi_buy_threshold):
                print("\n检测到买入信号组合:")
                if value_buy_signal:
                    print(f"基本面买入信号: PE偏离={self.pe_deviation:.2%}, PB偏离={self.pb_deviation:.2%}")
                print(f"技术面买入信号: MA5/MA10交叉, RSI={self.rsi_value:.2f}")
                
                # 计算仓位
                pos = self._calculate_position(bar, value_buy_signal, tech_buy_signal)
                if pos > 0:
                    price = bar.close_price
                    self.buy(price, pos)
                    print(f"发出买入委托: 价格={price:.2f}, 数量={pos}")
            else:
                print("\n不满足买入条件组合:")
                if self.current_pe > 0 and self.industry_pe > 0 and self.current_pb > 0 and self.industry_pb > 0:
                    print(f"基本面条件: PE偏离={self.pe_deviation:.2%}, PB偏离={self.pb_deviation:.2%}")
                print(f"技术面条件: 快线={self.fast_ma:.2f}, 慢线={self.slow_ma:.2f}, RSI={self.rsi_value:.2f}")
        
        # 检查卖出条件
        elif self.pos > 0:
            # 基本面卖出条件：PE或PB高于行业均值
            value_sell_signal = False
            if self.current_pe > 0 and self.industry_pe > 0 and self.current_pb > 0 and self.industry_pb > 0:
                value_sell_signal = (self.pe_deviation >= self.pe_high_threshold or 
                                    self.pb_deviation >= self.pb_high_threshold)
            
            # 技术面卖出条件：死叉或RSI超买
            tech_sell_signal = (self.fast_ma < self.slow_ma or 
                               self.rsi_value >= self.rsi_sell_threshold)
            
            # 综合卖出信号
            if value_sell_signal or tech_sell_signal:
                print("\n检测到卖出信号:")
                if value_sell_signal:
                    print(f"基本面卖出信号: PE偏离={self.pe_deviation:.2%}, PB偏离={self.pb_deviation:.2%}")
                if tech_sell_signal:
                    print(f"技术面卖出信号: MA5/MA10关系={self.fast_ma:.2f}/{self.slow_ma:.2f}, RSI={self.rsi_value:.2f}")
                
                price = bar.close_price
                self.sell(price, self.pos)
                print(f"发出卖出委托: 价格={price:.2f}, 数量={self.pos}")
            else:
                print("\n不满足卖出条件:")
                if self.current_pe > 0 and self.industry_pe > 0 and self.current_pb > 0 and self.industry_pb > 0:
                    print(f"基本面条件: PE偏离={self.pe_deviation:.2%}, PB偏离={self.pb_deviation:.2%}")
                print(f"技术面条件: 快线={self.fast_ma:.2f}, 慢线={self.slow_ma:.2f}, RSI={self.rsi_value:.2f}")
    
    def _calculate_position(self, bar: BarData, value_signal=False, tech_signal=False) -> int:
        """
        增强版仓位计算 - 根据信号强度动态调整仓位
        """
        if bar.close_price <= 0:
            return 0
        
        # 基础仓位
        base_position = self.fixed_size
        position_factor = 1.0
        
        # 根据基本面信号强度调整仓位
        if value_signal and self.pe_deviation <= self.pe_low_threshold and self.pb_deviation <= self.pb_low_threshold:
            # 两个指标都满足时增加仓位
            avg_deviation = (abs(self.pe_deviation) + abs(self.pb_deviation)) / 2
            position_factor += min(avg_deviation * 10, 1.0)  # 最多增加100%
        
        # 根据技术指标调整仓位
        if tech_signal:
            # RSI接近超卖区域时增加仓位
            if self.rsi_value < 35:
                position_factor += (35 - self.rsi_value) / 35
            
            # 强势上涨时增加仓位
            ma_ratio = self.fast_ma / self.slow_ma if self.slow_ma > 0 else 1.0
            if ma_ratio > 1.05:  # 快线比慢线高5%以上
                position_factor += min((ma_ratio - 1.0) * 5, 0.5)  # 最多增加50%
        
        # 计算最终仓位
        adjusted_position = int(base_position * position_factor)
        final_position = min(adjusted_position, self.max_pos)
        
        print(f"仓位计算: 基础={base_position}, 因子={position_factor:.2f}, 调整后={adjusted_position}, 最终={final_position}")
        return final_position

    def on_order(self, order: OrderData):
        """订单状态更新"""
        print(f"订单更新: ID={order.vt_orderid}, 方向={order.direction}, 价格={order.price}, "
              f"数量={order.volume}, 已成交={order.traded}, 状态={order.status}")
        
    def on_trade(self, trade: TradeData):
        """成交回报"""
        print(f"成交回报: ID={trade.vt_tradeid}, 方向={trade.direction}, "
              f"价格={trade.price}, 数量={trade.volume}")
        
        if trade.direction == Direction.LONG:
            self.entry_price = trade.price
            self.high_since_entry = trade.price
        
        self.put_event()
        
    def on_stop_order(self, stop_order: StopOrder):
        """条件单更新"""
        print(f"条件单更新: ID={stop_order.stop_orderid}, 状态={stop_order.status}")